import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain.agents import Tool
from langchain.agents import AgentType
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI
from langchain.utilities import SerpAPIWrapper
from langchain.agents import initialize_agent
from langchain.chains import LLMMathChain
from langchain.prompts import  MessagesPlaceholder


# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
api_base = os.getenv("OPENAI_API_BASE")
serpapi_api_key = os.getenv("SERPAPI_API_KEY")
os.environ["OPENAI_API_KEY"] = api_key
os.environ["OPENAI_API_BASE"] = api_base
os.environ["SERPAPI_API_KEY"] = serpapi_api_key

llm=ChatOpenAI(
    temperature=0,
    model="gpt-3.5-turbo",
)

# 记忆组件
memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True
)

# 构建一个搜索工具
search = SerpAPIWrapper()
# 创建一个数学计算工具
llm_math_chain = LLMMathChain(
    llm=llm,
    verbose=True
)
tools = [
    Tool(
        name = "Search",
        func=search.run,
        description="useful for when you need to answer questions about current events or the current state of the world"
    ),
    Tool(
        name="Calculator",
        func=llm_math_chain.run,
        description="useful for when you need to answer questions about math"
    ),
]

# 定义Agent
agent_chain = initialize_agent(
    tools,
    llm,
    agent=AgentType.OPENAI_FUNCTIONS, # 使用openai模型
    verbose=True, # 开启日志
    handle_parsing_errors=True,# 处理解析错误
    # 使用agent_kwargs传递参数，将chat_history传入
    agent_kwargs={
        "extra_prompt_messages": [MessagesPlaceholder(variable_name="chat_history"),
                                  MessagesPlaceholder(variable_name="agent_scratchpad")],
    },
    memory=memory # 记忆组件
)

agent_chain.run("你好，我叫Aceks，你呢？")
agent_chain.run("你会写代码吗？还记得我叫什么名字吗？")